import torch
from torch import nn


# 定义每批有多少个句子
batch_size = 5

# 定义每个句子中包含的词语数量
seq_len = 3

# 定义字典中词语的数量
total_word = 100

# 定义每个词语用数字向量表示的维度
embedding_dim = 10

# 定义隐藏层层神经元（LSTM）的数量
hidden_size = 5

# 定义隐藏层的数量
num_layer = 2

# 1. 准备数据
input = torch.randint(low=0, high=6, size=[batch_size, seq_len])

# 2. 实例化embedding，LSTM， GRU对象
embedding = nn.Embedding(total_word, embedding_dim)
lstm = nn.LSTM(input_size=embedding_dim, hidden_size=hidden_size, num_layers=num_layer, batch_first=True)
rev_lstm = nn.LSTM(input_size=embedding_dim, hidden_size=hidden_size, num_layers=num_layer, bidirectional=True)
gru = nn.GRU(input_size=embedding_dim, hidden_size=hidden_size, num_layers=num_layer)

# 3. 对数据进行embedding操作
embed = embedding(input)  # embed: [batch_size, seq_len, embedding_dim]

# 4. 将输出结果的形状修改为[seq_len, batch_size, embedding_dim]，因为此文件中的LSTM没有设置batch_first=True
# embed = embed.permute(1, 0, 2)

# 5. 初始化输入状态和神经元状态
# 单项LSTM
h_0 = torch.rand(num_layer, batch_size, hidden_size)
c_0 = torch.rand(num_layer, batch_size, hidden_size)
output, (h_n, c_n) = lstm(embed, (h_0, c_0))
# # 双向LSTM
# rev_h_0 = torch.rand(num_layer*2, batch_size, hidden_size)
# rev_c_0 = torch.rand(num_layer*2, batch_size, hidden_size)
# rev_output, (rev_h_n, rev_c_n) = rev_lstm(embed, (rev_h_0, rev_c_0))
# # GRU
# gru_h_0 = torch.rand(num_layer, batch_size, hidden_size)
# gru_output, gru_h_n = gru(embed, gru_h_0)

# # 双向LSTM
# print("双向LSTM")
# print(rev_output.size())
# print(rev_h_n.size())
# print(rev_c_n.size())
# # 正向：我在__， 只能预测句子的最后的空，所以是只能是正向的最后一个词
# # 在输出层取正向第一个句子（反向最后一个句子）的最后一个词（反向第一个词）的最后的输出
# # 第二个维度正向存储顺序是0->17, 反向存储是18->35
# rev_a = rev_output[-1, :1, :18]
# # print(a)
# # 在隐藏层取正向第一个句子最后的输出，rev_h_n的第0个维度是：第一层正向、第一层反向、第二层正向、第二层反向...，所以要取-2
# rev_b = rev_h_n[-2, :1, :]
# print(rev_b == rev_a)
# print("双向LSTM结束")
# print("*" * 50)
#
# # 反向：__在家，只能预测第一个空，所以只能是反向的最后一个词（正向的第一个词）
# # 在输出层取反向最后一个句子（正向第一个句子）的最后一个词（正向第一个词）的最后的输出
# rev_a = rev_output[0, :1, 18:]
# # 在隐藏层取反向最后一个句子（正向第一个句子）最后的输出
# rev_b = rev_h_n[-1, :1, :]
# print(rev_a == rev_b)
#
# print("*" * 50)
#
# # 在输出层取反向第一个句子（正向最后一个句子）的最后一个词（正向第一个词）的最后的输出
# rev_a2 = rev_output[0, -1, 18:]
# # 在隐藏层取反向第一个句子的最后的输出
# rev_b2 = rev_h_n[-1, -1, :]
# print(rev_a2 == rev_b2)
#
# print("*" * 100)

# 单项LSTM  output: [seq_len, batch_size, hidden_size]， h_n: [num_layers, batch_size, hidden_size]
print("这里是LSTM")
print("output:", output)
print("h_n:", h_n)
print(output.size())
print(h_n.size())
print(c_n.size())
# 在输出层取最后一个句子的最后一个词的最后的输出
a = output[-1, -1, :]
# 在隐藏层取最后一个句子的最后的输出
b = h_n.permute(1, 0, 2)
b = h_n[-1, -1, :]
print(a == b)
print("LSTM结束")
print("*" * 100)
print(output[-2, -1, :] == h_n[-1, -2, :])
"""
隐藏层只存储了每个句子的最后一个词的最后的输出
输出层存储了每个句子的每个词的最后的输出
"""

# # GRU  和单项LSTM效果相同
# print("GRU开始")
# print(gru_output.size())
# print(gru_h_n.size())
# # 在输出层取第一个句子的最后一个词的输出
# gru_a = gru_output[-1, 0, :]
# # 取隐藏层最后一个句子最后输出
# gru_b = gru_h_n[-1, 0, :]
# print(gru_a == gru_b)
# print("GRU结束")



